#lang racket

(provide (all-defined-out))

(define enable-cps #t)
(define commetted  #f)

(define wordsize       4)
(define blocksize      12)

(define compile-port
  (make-parameter
   (current-output-port)
   (lambda (p)
     (unless (output-port? p)
       (error 'compile-port (format "Not an output port ~s." p)))
     p)))

(define (emit . args)
  (apply fprintf (compile-port) args)
  (newline (compile-port)))

(define (emit-comment-start emitter)
  (if commetted (emit "#####  ~s: start  ######" emitter)
      '()))
(define (emit-comment-end emitter)
  (if commetted (emit "#####  ~s: end    ######" emitter)
      '()))

(define temp-registers
  '($ra $t0 $t1 $t2 $t3 $t4 $t5 $t6 $t7))

(define save-registers
  '($s0 $s1 $s2 $s3 $s4 $s5 $s6 $s7))

(define (preserve-registers cmd)
  (let loop ([regs temp-registers] [count 0])
    (unless (null? regs)
      (let ([reg (first regs)])
        (if (member reg temp-registers)
          (cmd reg (* count wordsize))
          '())
        (loop (rest regs) (+ count 1))))))

(define (backup-registers)
  (preserve-registers
   (lambda (name num)
     (emit "  sw ~a, ~s($sp)" name num))))

(define (restore-registers)
  (preserve-registers
   (lambda (name num)
     (emit "  lw ~a, ~s($sp)" name num))))

(define type-nil -1)
(define type-eof -2)
(define type-pair 256)
(define type-string 0)
(define type-symbol 1)
(define type-boolean 2)
(define type-integer 3)
(define type-fraction 4)
(define type-decimal 5)
(define type-character 6)
(define type-vector 7)
(define type-procedure 8)

(define bool-f         0)
(define bool-t         1)

(define (immediate? x)
  (or (fixnum? x) (boolean? x) (null? x) (char? x)))

(define (emit-immediate-nil)
  (emit "  li   $s1, -1")
  (emit "  li   $s2, 0")
  (emit "  li   $s3, 0"))

(define (emit-immediate-fixnum x)
  (emit "  li   $s1, 3")
  (emit "  li   $s2, ~s" x)
  (emit "  li   $s3, 1"))

(define (emit-immediate-boolean x)
  (emit "  li   $s1, 2")
  (emit "  li   $s2, ~s" (if x 1 0))
  (emit "  li   $s3, 0"))

(define (emit-immediate-char x)
  (emit "  li   $s1, 6")
  (emit "  li   $s2, ~s" (char->integer x))
  (emit "  li   $s3, 0"))

(define (emit-immediate x)
  (emit-comment-start 'emit-immediate)
  (cond
   [(fixnum? x) (emit-immediate-fixnum x)]
   [(boolean? x) (emit-immediate-boolean x)]
   [(null? x) (emit-immediate-nil)]
   [(char? x) (emit-immediate-char x)]
   [else #f])
  (emit-comment-end 'emit-immediate))


(define (emit-stack-save si)
  (emit "  sw $s1, ~s($sp)     " si)
  (emit "  sw $s2, ~s($sp)     " (- si wordsize))
  (emit "  sw $s3, ~s($sp)     " (- si wordsize wordsize)))

(define (emit-stack-load si)
  (emit "  lw $s1, ~s($sp)     " si)
  (emit "  lw $s2, ~s($sp)     " (- si wordsize))
  (emit "  lw $s3, ~s($sp)     " (- si wordsize wordsize)))

(define (next-stack-index si)
  (- si blocksize))

(define (emit-cmp-type type)
  (emit "  move  $a0, $s1               ")
  (emit "  li    $t0, ~s                " type)
  (emit-cmp-bool))

(define (emit-cmp-type-and-value type value)
  (emit "  move  $t1, $s1               ")
  (emit "  move  $t2, $s2               ")
  (emit "  move  $a0, $t1               ")
  (emit "  li    $t0, ~s                " type)
  (emit-cmp-bool)
  (emit "  move  $t1, $s2               ")
  (emit "  move  $a0, $t2               ")
  (emit "  li    $t0, ~s                " value)
  (emit-cmp-bool)
  (emit "  move  $t2, $s2               ")
  (emit "  and   $s2, $t1, $t2          "))

(define (emit-cmp-bool . args)
  (let ((label1 (unique-label))
        (label2 (unique-label)))
    (emit-comment-start 'emit-cmp-bool)
    (emit "  ~s    $a0, $t0, ~a  # if $a0 == ?    " (if (null? args) 'beq (car args)) label1)
    (emit "  li    $v0, ~s       # $v0 = 0        " bool-f)
    (emit "  j     ~a            #                " label2)
    (emit "~a:                                    " label1)
    (emit "  li    $v0, ~s       # $v0 = 1        " bool-t)
    (emit "~a:                                    " label2)
    (emit "  li    $s1, 2        #                ")
    (emit "  move  $s2, $v0      #                ")
    (emit "  li    $s3, 0        #                ")
    (emit-comment-end 'emit-cmp-bool)))

(define (emit-binop si env arg1 arg2)
  (emit-comment-start 'emit-binop)
  ;(emit-expr si env arg1)
  ;(emit "  move  $t1, $s2      # $t1 = $s2      ")
  ;(emit-expr si env arg2)
  ;(emit "  move  $t2, $s2      # $t2 = $s2      ")
  (emit-expr si env arg1)
  (emit-stack-save si)
  (emit-expr (next-stack-index si) env arg2)
  (emit-comment-end 'emit-binop))

(define (emit-cmp-binop setx si env arg1 arg2)
  (emit-comment-start 'emit-cmp-binop)
  (emit-binop si env arg1 arg2)
  (emit "  lw $t1, ~s($sp)     " (- si wordsize))
  (emit "  move  $t2, $s2            ")
  (emit "  move  $a0, $t1      # $a0 = $t1      ")
  (emit "  move  $t0, $t2      # $t0 = $t2      ")
  (emit-cmp-bool setx)
  (emit-comment-end 'emit-cmp-binop))

(define (primitive-label name)
  (let ([lst (map (lambda (c)
		    (case c
		      [(#\-) #\_]
		      [(#\!) #\b]
		      [(#\=) #\e]
		      [(#\>) #\g]
		      [(#\?) #\p]
		      [else c]))
		  (string->list (symbol->string name)))])
  (string->symbol (format "P_~a" (list->string lst)))))


(define primitives
  `(fxadd1
    ,(lambda (si env arg)
       (emit-comment-start 'fxadd1)
       (emit-expr si env arg)
       (emit "  add   $s2, $s2, 1   # fxadd1 ")
       (emit-comment-end 'fxadd1))
    fxsub1
    ,(lambda (si env arg)
       (emit-comment-start 'fxsub1)
       (emit-expr si env arg)
       (emit "  add   $s2, $s2, -1  # fxsub1 ")
       (emit-comment-end 'fxsub1))
    fixnum->char
    ,(lambda (si env arg)
       (emit-comment-start 'fixnum->char)
       (emit-expr si env arg)
       (emit "  li    $s1, 6  # fixnum->char ")
       (emit "  li    $s3, 0                 ")
       (emit-comment-end 'fixnum->char))
    char->fixnum
    ,(lambda (si env arg)
       (emit-comment-start 'char->fixnum)
       (emit-expr si env arg)
       (emit "  li    $s1, 3  # char->fixnum ")
       (emit "  li    $s3, 1                 ")
       (emit-comment-end 'char->fixnum))
    fixnum?
    ,(lambda (si env arg)
       (emit-comment-start 'fixnum?)
       (emit-expr si env arg)
       (emit-cmp-type type-integer)
       (emit-comment-end 'fixnum?))
    fxzero?
    ,(lambda (si env arg)
       (emit-comment-start 'fxzero?)
       (emit-expr si env arg)
       (emit-cmp-type-and-value type-integer 0)
       (emit-comment-end 'fxzero?))
    null?
    ,(lambda (si env arg)
       (emit-comment-start 'null?)
       (emit-expr si env arg)
       (emit-cmp-type type-nil)
       (emit-comment-end 'null?))
    eof-object?
    ,(lambda (si env arg)
       (emit-comment-start 'eof-object?)
       (emit-expr si env arg)
       (emit-cmp-type type-eof)
       (emit-comment-end 'eof-object?))
    eof-object
    ,(lambda (si env)
       (emit-comment-start 'eof-object)
       (emit "  li    $s1, -2 # eof          ")
       (emit "  li    $s2, 0                 ")
       (emit "  li    $s3, 0                 ")
       (emit-comment-end 'eof-object))
    boolean?
    ,(lambda (si env arg)
       (emit-comment-start 'boolean?)
       (emit-expr si env arg)
       (emit-cmp-type type-boolean)
       (emit-comment-end 'boolean?))
    char?
    ,(lambda (si env arg)
       (emit-comment-start 'char?)
       (emit-expr si env arg)
       (emit-cmp-type type-character)
       (emit-comment-end 'char?))
    not
    ,(lambda (si env arg)
       (emit-comment-start 'not)
       (emit-expr si env arg)
       (emit-cmp-type-and-value type-boolean 0)
       (emit-comment-end 'not))
    fxlognot
    ,(lambda (si env arg)
      (emit-comment-start 'fxlognot)
      (emit-expr si env arg)
      (emit "  not   $s2, $s2                ")
      (emit-comment-end 'fxlognot))
    fx+
    ,(lambda (si env arg1 arg2)
       (emit-comment-start 'fx+)
       (emit-binop si env arg1 arg2)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  add   $s2, $t1, $s2 ")
       (emit-comment-end 'fx+))
    fx-
    ,(lambda (si env arg1 arg2)
       (emit-comment-start 'fx-)
       (emit-binop si env arg1 arg2)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  sub   $s2, $t1, $s2 ")
       (emit-comment-end 'fx-))
    fx*
    ,(lambda (si env arg1 arg2)
       (emit-comment-start 'fx*)
       (emit-binop si env arg1 arg2)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  mul   $s2, $t1, $s2 ")
       (emit-comment-end 'fx*))
    $fxquotient
    ,(lambda (si env arg1 arg2)
       (emit-comment-start '$fxquotient)
       (emit-binop si env arg1 arg2)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  move  $t2, $s2")
       (emit "  div   $t1, $t2")
       (emit "  mfhi  $s2")
       (emit-comment-end '$fxquotient))
    $fxremainder
    ,(lambda (si env arg1 arg2)
       (emit-comment-start '$fxremainder)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  move  $t2, $s2")
       (emit "  div   $t1, $t2")
       (emit "  mflo  $s2")
       (emit-comment-end '$fxremainder))
    fxlogor
    ,(lambda (si env arg1 arg2)
       (emit-comment-start 'fxlogor)
       (emit-binop si env arg1 arg2)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  or    $s2, $t1, $s2 ")
       (emit-comment-end 'fxlogor))
    fxlogand
    ,(lambda (si env arg1 arg2)
       (emit-comment-start 'fxlogand)
       (emit-binop si env arg1 arg2)
       (emit "  lw $t1, ~s($sp)     " (- si wordsize))
       (emit "  and   $s2, $t1, $s2 ")
       (emit-comment-end 'fxlogand))
    fx=
    ,(lambda (si env arg1 arg2)
      (emit-comment-start 'fx=)
      (emit-binop si env arg1 arg2)
      (emit-cmp-binop 'beq si env arg1 arg2)
      (emit-comment-end 'fx=))
    fx<
    ,(lambda (si env arg1 arg2)
      (emit-comment-start 'fx<)
      (emit-binop si env arg1 arg2)
      (emit-cmp-binop 'blt si env arg1 arg2)
      (emit-comment-end 'fx<))
    fx<=
    ,(lambda (si env arg1 arg2)
      (emit-comment-start 'fx<=)
      (emit-binop si env arg1 arg2)
      (emit-cmp-binop 'ble si env arg1 arg2)
      (emit-comment-end 'fx<=))
    fx>
    ,(lambda (si env arg1 arg2)
      (emit-comment-start 'fx>)
      (emit-binop si env arg1 arg2)
      (emit-cmp-binop 'bgt si env arg1 arg2)
      (emit-comment-end 'fx>))
    fx>=
    ,(lambda (si env arg1 arg2)
       (emit-comment-start 'fx>=)
       (emit-binop si env arg1 arg2)
       (emit-cmp-binop 'bge si env arg1 arg2)
       (emit-comment-end 'fx>=))
    cons
    ,(lambda (si env arg1 arg2)
       (emit-binop si env arg1 arg2)
       (emit-stack-save (next-stack-index si))
       (emit-heap-alloc (* 2 blocksize))
       (emit-stack-to-heap si 0)
       (emit "  move $t0, $v0  ")
       (emit-stack-to-heap (next-stack-index si) blocksize)
       (emit "  li   $s1, ~s   " type-pair)
       (emit "  move $s2, $t0  ")
       (emit "  addi $s3, $v0, ~s  " blocksize))
    pair?
    ,(lambda (si env arg)
       (emit-comment-start 'pair?)
       (emit-expr si env arg)
       (emit-cmp-type type-pair)
       (emit-comment-end 'pair?))
    car
    ,(lambda (si env arg)
       (emit-expr si env arg)
       (emit "  move $v0, $s2  ")
       (emit "  lw $s1, 0($v0)     ")
       (emit "  lw $s2, 4($v0)     ")
       (emit "  lw $s3, 8($v0)     "))
    cdr
    ,(lambda (si env arg)
       (emit-expr si env arg)
       (emit "  move $v0, $s3  ")
       (emit "  lw $s1, 0($v0)     ")
       (emit "  lw $s2, 4($v0)     ")
       (emit "  lw $s3, 8($v0)     "))
    set-car!
    ,(lambda (si env cell val)
       (emit-binop si env val cell)
       (emit "  move $v0, $s2  ")
       (emit-stack-load si)
       (emit "  sw $s1, 0($v0)     ")
       (emit "  sw $s2, 4($v0)     ")
       (emit "  sw $s3, 8($v0)     "))
    set-cdr!
    ,(lambda (si env cell val)
       (emit-binop si env val cell)
       (emit "  move $v0, $s3  ")
       (emit-stack-load si)
       (emit "  sw $s1, 0($v0)     ")
       (emit "  sw $s2, 4($v0)     ")
       (emit "  sw $s3, 8($v0)     "))
    eq?
    ,(lambda (si env arg1 arg2)
       (let ((label1 (unique-label))
             (label2 (unique-label)))
         (emit-binop si env arg1 arg2)
         (emit "  lw $t0, ~s($sp)     " si)
         (emit "  lw $t1, ~s($sp)     " (- si wordsize))
         (emit "  lw $t2, ~s($sp)     " (- si wordsize wordsize))
         (emit "  xor   $t0, $s1, $t0")
         (emit "  xor   $t1, $s2, $t1")
         (emit "  xor   $t2, $s3, $t2")
         (emit "  or    $a0, $t0, $t1")
         (emit "  or    $a0, $a0, $t2")
         (emit "  beq   $a0, $zero, ~a                 " label1)
         (emit "  li    $v0, ~s       # $v0 = 0        " bool-f)
         (emit "  j     ~a            #                " label2)
         (emit "~a:                                    " label1)
         (emit "  li    $v0, ~s       # $v0 = 1        " bool-t)
         (emit "~a:                                    " label2)
         (emit "  li    $s1, 2        #                ")
         (emit "  move  $s2, $v0      #                ")
         (emit "  li    $s3, 0        #                ")))

    make-vector
    ,(lambda (si env length)
       (emit-expr si env length)
       (emit-stack-save si)
       (emit-heap-alloc-dynamic)
       (emit "  li    $s1, ~s      " type-vector)
       (emit "  move  $s2, $v0     ")
       (emit "  lw    $s3, ~s($sp) " (- si wordsize)))
    vector?
    ,(lambda (si env arg)
       (emit-comment-start 'vector?)
       (emit-expr si env arg)
       (emit-cmp-type type-vector)
       (emit-comment-end 'vector?))
    vector-length
    ,(lambda (si env arg)
       (emit-comment-start 'vector?)
       (emit-expr si env arg)
       (emit "  li    $s1, ~s      " type-integer)
       (emit "  move  $s2, $s3     ")
       (emit "  li    $s3, 0       ")
       (emit-comment-end 'vector?))
    vector-set!
    ,(lambda (si env vector index value)
       (emit-expr si env index)
       (emit-stack-save si)
       (emit-expr (next-stack-index si) env value)
       (emit-stack-save (next-stack-index si))
       (emit-expr si env vector)
       (emit "  lw    $t1, ~s($sp)    " (- si wordsize))
       (emit "  sll   $t2, $t1, 3     ")
       (emit "  sll   $t3, $t1, 2     ")
       (emit "  add   $t1, $t2, $t3   ")
       (emit "  add   $v0, $s2, $t1   ")
       (emit-stack-to-heap (next-stack-index si) 0))
    vector-ref
    ,(lambda (si env vector index)
       (emit-expr si env index)
       (emit-stack-save si)
       (emit-expr si env vector)
       (emit "  lw    $t1, ~s($sp)    " (- si wordsize))
       (emit "  sll   $t2, $t1, 3     ")
       (emit "  sll   $t3, $t1, 2     ")
       (emit "  add   $t1, $t2, $t3   ")
       (emit "  add   $v0, $s2, $t1   ")
       (emit-heap-load))

    make-string
    ,(lambda (si env length)
       (emit-expr si env length)
       (emit-stack-save si)
       (emit "  addi $a0, $s2, 1")
       (emit "  li   $v0, 9")
       (emit "  syscall")
       (emit "  li    $s1, ~s      " type-string)
       (emit "  move  $s2, $v0     ")
       (emit "  lw    $s3, ~s($sp) " (- si wordsize)))
    string?
    ,(lambda (si env arg)
       (emit-expr si env arg)
       (emit-cmp-type type-string))
    string-length
    ,(lambda (si env arg)
       (emit-expr si env arg)
       (emit "  li    $s1, ~s      " type-integer)
       (emit "  move  $s2, $s3     ")
       (emit "  li    $s3, 0       "))
    string-set!
    ,(lambda (si env vector index value)
       (emit-expr si env index)
       (emit-stack-save si)
       (emit-expr (next-stack-index si) env value)
       (emit-stack-save (next-stack-index si))
       (emit-expr si env vector)
       (emit "  lw    $t1, ~s($sp)    " (- si wordsize))
       (emit "  add   $v0, $s2, $t1   ")
       (emit-stack-load (next-stack-index si))
       (emit "  sb $s2, 0($v0)      "))
    string-ref
    ,(lambda (si env vector index)
       (emit-expr si env index)
       (emit-stack-save si)
       (emit-expr si env vector)
       (emit "  lw    $t1, ~s($sp)    " (- si wordsize))
       (emit "  add   $v0, $s2, $t1   ")
       (emit "  li    $s1, ~s         " type-character)
       (emit "  lbu   $s2, 0($v0)     ")
       (emit "  li    $s3, 0          "))
    char=
    ,(lambda (si env arg1 arg2)
       (emit-cmp-binop 'beq si env arg1 arg2))
    
    constant-ref
    ,(lambda (si env constant)
      (emit-comment-start 'constant-ref)
      (emit "  lw    $a0, ~s       # $a0 = ?       " constant)
      (emit-comment-end 'constant-ref))
    constant-init
    ,(lambda (si env constant value)
      (emit-comment-start 'constant-init)
      (emit ".data  ~s    .word 0       # $a0 = ?  " constant)
      (emit-expr si env value)
      (emit "  sw    $a0, ~s       # $a0 = ?       " constant)
      (emit-comment-end 'constant-init))
    primitive-ref
    ,(lambda (si env prim-name)
      (let ([label (primitive-label prim-name)]
            [done-label (unique-label)])
        (emit-comment-start 'primitive-ref)
        (emit "  lw    $a0, ~s       # $a0 = ?       " label)
        (emit-comment-end 'primitive-ref)))))

(define (emit-heap-alloc-dynamic)
  (emit "  sll   $a0, $s2, 3     ")
  (emit "  sll   $a1, $s2, 2     ")
  (emit "  add   $a0, $a0, $a1   ")
  (emit "  li   $v0, 9")
  (emit "  syscall"))

(define (emit-heap-alloc size)
  (emit "  li   $a0, ~a" size)
  (emit "  li   $v0, 9")
  (emit "  syscall"))
(define (emit-stack-to-heap si offset)
  (emit-stack-load si)
  (emit "  sw $s1, ~s($v0)     " offset)
  (emit "  sw $s2, ~s($v0)     " (+ offset wordsize))
  (emit "  sw $s3, ~s($v0)     " (+ offset wordsize wordsize)))

;(define (emit-heap-save offset)
;  (emit "  move $a0, $v0  ")
;  (emit "  lw $s1, ~s($a0)     " 0)
;  (emit "  lw $s2, ~s($a0)     " -4)
;  (emit "  lw $s3, ~s($a0)     " -8))
(define (emit-heap-load)
  (emit "  lw $s1, 0($v0)     ")
  (emit "  lw $s2, 4($v0)     ")
  (emit "  lw $s3, 8($v0)     "))


(define (is-prim? x lst)
  (cond
    ((null? lst) #f)
    ((eq? x (car lst)) #t)
    (else (is-prim? x (cddr lst)))))

(define (is-lib-prim? x lst) (is-prim? x lst))

(define (primitive? x)
  (and (symbol? x) (is-prim? x primitives)))

(define (primcall? expr)
  (and (pair? expr) (primitive? (car expr))))

(define (primitive-emitter x)
  (or (code x primitives) (error 'primitive-emitter (format "primitive ~s has no emitter" x))))

(define (code x lst)
  (cond
    ((null? lst) #f)
    ((eq? x (car lst)) (cadr lst))
    (else (code x (cddr lst)))))

(define (emit-primcall si env expr)
  (let ([prim (car expr)] [args (cdr expr)])
    (apply (primitive-emitter prim) si env args)))

(define unique-label
  (let ([count 0])
    (lambda ()
      (let ([L (string->symbol (format "L_~s" count))])
        (set! count (add1 count))
        L))))

(define (if? expr)
  (and (tagged-list 'if expr)
       (or (= 3 (length (cdr expr)))
           (error 'if? "malformed if ~s" expr))))
(define if-test cadr)
(define if-conseq caddr)
(define if-altern cadddr)

(define (emit-label label)
  (emit "~a:" label))

(define (emit-call label)
  (emit "  jal ~a" label))

(define (emit-if si env tail expr)
  (let ([alt-label (unique-label)]
        [end-label (unique-label)])
    (emit-comment-start 'if)
    (emit-expr si env (if-test expr))
    (emit "  xori  $t0, $s1, ~s" type-boolean)
    (emit "  xori  $t1, $s2, ~s" bool-f)
    (emit "  or    $t2, $t0, $t1")
    (emit "  beq   $t2, $zero, ~a" alt-label)
    (emit-any-expr si env tail (if-conseq expr))
    (if (not tail) (emit "  j  ~a" end-label) '())
    (emit-label alt-label)
    (emit-any-expr si env tail (if-altern expr))
    (emit-label end-label)
    (emit-comment-end 'if)))

(define variable? symbol?)
(define (tagged-list tag expr)
  (and (list? expr) (not (null? expr)) (eq? (car expr) tag)))

(define (make-begin seq) (cons 'begin seq))
(define (begin? expr)
  (and (tagged-list 'begin expr)
       (or (not (null? (begin-seq expr)))
           (error 'begin? (format "empty begin")))))
(define begin-seq cdr)
(define (emit-begin si env tail expr)
  (emit-seq si env tail (begin-seq expr)))
(define (emit-seq si env tail seq)
  (cond
   [(null? seq) (error 'emit-seq "empty seq")]
   [(null? (rest seq)) (emit-any-expr si env tail (first seq))]
   [else
    (emit-expr si env (first seq))
    (emit-seq si env tail (rest seq))]))

(define make-let list)
(define (let-form? let-kind expr)
  (and (tagged-list let-kind expr)
       (or (not (null? (cddr expr)))
           (error 'let-form? (format "let without body ~s" expr)))))
(define let-kind car)
(define (any-let? expr)
  (and (pair? expr)
       (member (let-kind expr) '(let let* letrec))
       (let-form? (let-kind expr) expr)))
(define (let? expr) (let-form? 'let expr))
(define (let*? expr) (let-form? 'let* expr))
(define (letrec? expr) (let-form? 'letrec expr))
(define let-bindings cadr)
(define letrec-bindings let-bindings)
(define labels-bindings let-bindings)
(define (make-body lst)
  (if (null? (cdr lst))
      (car lst)
      (make-begin lst)))
(define (let-body expr)
  (make-body (cddr expr)))
(define letrec-body let-body)
(define labels-body let-body)
(define empty? null?)
(define (bind lhs rhs)
  (check-variable lhs)
  (list lhs rhs))
(define first car)
(define rest cdr)
(define rhs cadr)
(define (lhs binding)
  (check-variable (car binding)))
(define (check-variable var)
  (if (and (variable? var) (not (special? var)))
      var
      (error 'lhs (format "~s is not a variable" var))))
(define (make-initial-env bindings)
  bindings)
(define (extend-env var si env)
  (cons (list var si) env))
(define (lookup var env)
  (cond
   [(assv var env) => cadr]
   [else #f]))

(define (emit-let si env tail expr)
  (define (process-let bindings si new-env)
    (cond
     [(empty? bindings)
      (emit-any-expr si new-env tail (let-body expr))]
     [else
      (let ([b (first bindings)])
        (emit-expr si (if (let*? expr) new-env env) (rhs b))
        (emit-stack-save si)
        (process-let (rest bindings)
           (next-stack-index si)
           (extend-env (lhs b) si new-env)))]))
  (process-let (let-bindings expr) si env))

(define (extend-env-with si env lvars k)
  (if (null? lvars)
      (k si env)
      (extend-env-with
       (next-stack-index si)
       (extend-env (first lvars) si env)
       (rest lvars)
       k)))

(define closuretag 0)

(define (free-var offset)
  (list 'free (- offset closuretag)))
(define (free-var? fv)
  (tagged-list 'free fv))
(define free-var-offset cadr)

(define (close-env-with offset env lvars k)
  (if (null? lvars)
      (k env)
      (close-env-with
       (+ offset wordsize)
       (extend-env (first lvars) (free-var offset) env)
       (rest lvars)
       k)))

(define label? symbol?)

(define (emit-variable-ref si env var)
  (cond
   [(lookup var env) =>
    (lambda (v)
      (cond 
       [(free-var? v)
        (emit "  mov ~s(%rdi), %rax" (free-var-offset v))]
       [(number? v)
        (emit-stack-load v)]
       [(label? v)
        (emit-closure si env (make-closure v '()))]
       [else (error 'emit-variable-ref (format "looked up unknown value ~s for var ~s" v var))]))]
   [else (error 'emit-variable-ref (format "undefined variable ~s" var))]))

(define (emit-ret-if tail)
  (if tail (emit-return) '()))

(define (emit-expr si env expr)
  (emit-any-expr si env #f expr))

(define (emit-tail-expr si env expr)
  (emit-any-expr si env #t expr))

(define (emit-any-expr si env tail expr)
  (cond
   [(immediate? expr) (emit-immediate expr) (emit-ret-if tail)]
   [(variable? expr) (emit-variable-ref env expr) (emit-ret-if tail)]
   [(closure? expr) (emit-closure si env expr) (emit-ret-if tail)]
   [(if? expr) (emit-if si env tail expr)]
   [(or (let? expr) (let*? expr)) (emit-let si env tail expr)]
   [(begin? expr) (emit-begin si env tail expr)]
   [(primcall? expr) (emit-primcall si env expr) (emit-ret-if tail)]
   [(app? expr env) (emit-app si env tail expr)]
   [else (error 'emit-expr (format "~s is not an expression" expr))]))

(define (closure-conversion expr)
  (let ([labels '()]
        [top-env '()]
        [top-procs '()])
    (define (transform-letrec expr)
      (let ([top-bindings
             (map (lambda (binding) (bind (lhs binding) (unique-label)))
                  (letrec-bindings expr))])
        (set! top-env (make-initial-env top-bindings))
        (set! top-procs (map lhs top-bindings))
        (for-each (lambda (lvar-lambda lvar-label)
                    (transform (rhs lvar-lambda) (rhs lvar-label)))
                  (letrec-bindings expr) top-bindings)
        (transform (letrec-body expr))))
    (define (transform expr . label)
      (cond
       [(lambda? expr)
        (let ([label (or (and (not (null? label)) (car label)) (unique-label))]
              [fvs (filter (lambda (v) (not (member v top-procs))) (free-vars expr))])
          (set! labels
                (cons (bind label
                            (make-code (lambda-formals expr)
                                       fvs
                                       (transform (lambda-body expr))))
                    labels))
          (make-closure label fvs))]
       [(any-let? expr)
        (make-let (let-kind expr)
                  (map (lambda (binding)
                         (bind (lhs binding) (transform (rhs binding))))
                       (let-bindings expr))
                  (transform (let-body expr)))]
       [(list? expr)
        (map transform expr)]
       [else
        expr]))
    (let* ([body (if (letrec? expr)
                     (transform-letrec expr)
                     (transform expr))])
      (make-top top-env (make-let 'labels labels body)))))

(define make-top list)
(define top-env car)
(define top-expr cadr)

(define (special? symbol)
  (or (member symbol '(if begin let let* letrec lambda closure))
      (primitive? symbol)))

(define (flatmap f . lst)
  (apply append (apply map f lst)))

(define (free-vars expr)
  (cond
   [(and (variable? expr) (not (special? expr))) (list expr)]
   [(lambda? expr) (filter (lambda (v) (not (member v (lambda-formals expr))))
                           (free-vars (lambda-body expr)))]
   [(let? expr)
    (append
     (flatmap free-vars (map rhs (let-bindings expr)))
     (filter (lambda (v) (not (member v (map lhs (let-bindings expr)))))
             (free-vars (let-body expr))))]
   [(let*? expr)
    (if (null? (let-bindings expr))
        (free-vars (let-body expr))
        (append
         (free-vars (rhs (first (let-bindings expr))))
         (filter (lambda (v) (not (eq? v (lhs (first (let-bindings expr))))))
                 (free-vars (make-let 'let* (rest (let-bindings expr)) (let-body expr))))))]
   [(list? expr) (flatmap free-vars expr)]
   [else '()]))

(define (emit-top top index)
  (emit-labels (top-expr top) (top-env top)) index)

(define (emit-labels expr env index)
  (let* ([bindings (labels-bindings expr)]
         [labels (map lhs bindings)]
         [codes (map rhs bindings)])
    (for-each (emit-code env) codes labels)
    (emit-scheme-entry-i (labels-body expr) env index)))

(define (lambda? expr) (tagged-list 'lambda expr))
(define lambda-formals cadr)
(define (lambda-body expr) (make-body (cddr expr)))

(define (make-closure label fvs)
  (cons 'closure (cons label fvs)))
(define (closure? expr) (tagged-list 'closure expr))
(define closure-label cadr)
(define closure-free-vars cddr)
(define (emit-closure si env expr)
  (let ([label (closure-label expr)]
        [fvs (closure-free-vars expr)])
    (emit-heap-alloc (* (add1 (length fvs)) blocksize))
    (emit "  la  $a0, ~s  #movq ?, (%rax)" label)
    (emit "  sw  $a0, ~s  #movq ?, (%rax)")
    (unless (null? fvs)
      (emit "  mov %rax, %rdx")
      (let loop ([fvs fvs] [count 1])
        (unless (null? fvs)
          (emit-variable-ref si env (first fvs))
          (emit "  mov %rax, ~s(%rdx)" (* count blocksize))
          (loop (rest fvs) (add1 count))))
      (emit "  mov %rdx, %rax"))
    (emit "  or $~s, %rax" closuretag)))

(define (make-code formals free body)
  (list 'code formals free body))
(define code-formals cadr)
(define code-free-variables caddr)
(define code-body cadddr)
(define (emit-code env)
  (lambda (expr label)
    (emit-function-header label)
    (let ([fmls (code-formals expr)]
          [fvs (code-free-variables expr)]
          [body (code-body expr)])
      (extend-env-with (- wordsize) env fmls (lambda (si env)
        (close-env-with wordsize env fvs (lambda (env)
          (emit-tail-expr si env body))))))))

(define (app? expr env)
  (and (list? expr) (not (null? expr))))
(define call-target car)
(define call-args cdr)

(define (emit-app si env tail expr)
  (define (emit-arguments si args)
    (unless (empty? args)
      (emit-expr si env (first args))
      (emit-stack-save si)
      (emit-arguments (- si blocksize) (rest args))))
  (define (move-arguments si delta args)
    (unless (or (= delta 0) (empty? args))
      (emit-stack-load si)
      (emit-stack-save (+ si delta))
      (move-arguments (- si blocksize) delta (rest args))))
  (let ([target-proc (proc (call-target expr) env)])
    (cond
     [(not tail)
      (emit-arguments (- si (* 2 wordsize)) (call-args expr))
      (when (not target-proc)
            (emit-expr si env (call-target expr))
            (emit "  mov %rdi, ~s(%rsp)" si)
            (emit "  mov %rax, %rdi")
            (emit-heap-load (- closuretag)))
      (emit-adjust-base si)
      (cond
       [target-proc => emit-call]
       [else (emit-call "*%rax")])
      (emit-adjust-base (- si))
      (when (not target-proc)
            (emit "  mov ~s(%rsp), %rdi" si))]
     [else ; tail
      (emit-arguments si (call-args expr))
      (when (not target-proc)
            (emit-expr si env (call-target expr))
            (emit "  mov %rax, %rdi"))
      (move-arguments si (- (+ si wordsize)) (call-args expr))
      (when (not target-proc)
            (emit "  mov %rdi, %rax")
            (emit-heap-load (- closuretag)))
      (cond
       [target-proc => emit-jmp]
       [else (emit-jmp "*%rax")])])))

(define (proc expr env)
  (cond 
   [(and (variable? expr) (lookup expr env)) =>
    (lambda (val) (and (label? val) val))]
   [else #f]))
    


(define (emit-adjust-base si)
  (unless (= si 0) (emit "  addi $sp, $sp, ~s" si)))

(define (emit-function-header f)
  (emit "  .text")
  (emit "  .globl ~a" f)
  (emit-label f))

(define (emit-save-return-address)
  (emit " addi  $sp, $sp, -24 ")
  (emit " sw $ra, 0($sp)"))

(define (emit-return)
  (emit " lw $ra, 0($sp)")
  (emit " addi  $sp, $sp, 24 ")
  (emit "  jr  $ra"))

(define (emit-jmp label)
  (emit "  j   ~a" label))

;;; test-program

(define (make-enumeration start end)
  (if (> start end) '()
      (if (= start end) (list start)
          (append (list start) (make-enumeration (+ start 1) end)))))

(define (emit-scheme-entry-i expr env index)
  (emit-function-header (string-append "L_scheme_entry_" index))
  (emit-save-return-address)
  (emit-tail-expr (- blocksize) env expr))

(define (emit-test-result index)
  (emit-call (string-append "L_scheme_entry_" index))
  (emit "  li   $a0, 12")
  (emit "  li   $v0, 9")
  (emit "  syscall")
  (emit "  move $a0, $v0")
  (emit "  sw   $s1, 0($a0)")
  (emit "  sw   $s2, 4($a0)")
  (emit "  sw   $s3, 8($a0)")
  (emit-call "print_sexp")
  (emit "  li   $a0, '\\n'")
  (emit "  li   $v0, 11")
  (emit "  syscall"))

(define (emit-test-program expr-list testname)
  (define expr-indexes (make-enumeration 1 (length expr-list)))
  (define expr-name-indexes (map (lambda (i) (string-append testname "_" (number->string i))) expr-indexes))
  (emit-function-header (string-append "main_" testname))
  (emit "  li   $a0, 0")
  (emit "  li   $v0, 9")
  (emit "  syscall")
  (emit "  move  $s0, $v0")
  (emit "  addi $sp, $sp, -36")
  (backup-registers)
  ;(emit "  mov %ecx, %esi")
  ;(emit "  mov 12(%esp), %ebp")
  ;(emit "  mov 8(%esp), %esp")
  ;(emit "  mov $0, %edi")
  (map (lambda (i) (emit-test-result i)) expr-name-indexes)
  ;(emit "  mov %esi, %ecx")
  (restore-registers)
  (emit "  addi $sp, $sp, 36")
  (emit "  li   $v0, 17  # exit")
  (emit "  syscall             ")

  ;  (cond 
  ;   [(letrec? program) (emit-letrec program)]
  ;   [else (emit-scheme-entry program (make-initial-env '() '()))])
  ;  (map (lambda (e i) (emit-scheme-entry-i e '() i)) expr-list expr-name-indexes)

  (map (lambda (e i) (emit-top (closure-conversion e) i)) expr-list expr-name-indexes))


